import torch, collections

print('notice !!!! ----> use python3 run this script!!! \n')
model_path = 'epoch76Loss0.068603Acc0.983750.pt'  # 要转换的pytorch的pt模型文件路径
OUT_ONNX = './LeNet.onnx'  # 要输出的onnx模型文件路径

device = 'cuda' if torch.cuda.is_available() else 'cpu'  # 判断能否加载gpu类型的tensor，以进行gpu并行运算
x = torch.randn(1, 1, 28, 28)  # 设置模型输入信息的shape
x = x.to(device)  # 设置tensor的数据类型，若device='cpu'，则进行tensor为原类型，使用cpu来运算；若device='cuda'，则转为cuda类型的tensor，使用gpu运算
input_names = ["input"]  # 每一层的输入的名称，如第一层输入将在模型中为input[0]，第二层为input[1]
out_names = ["output"]  # 每一层的输出名称，如第一层输入将在模型中为output[0]，第二层为output[1]

xmodel = torch.load(model_path, map_location=device)  # 加载模型（map_location将根据device进行cpu或gpu的部署），此加载模型方式会加载模型的结构和参数
# model.load_state_dict(torch.load(model_path, map_location=device)['model']) # 此加载模型的方式只加载模型的参数


'''关于model.eval():
不启用 BatchNormalization 和 Dropout，保证BN和dropout不发生变化
pytorch框架会自动把BN和Dropout固定住，不会取平均，而是用训练好的值
不然的话，一旦test的batch_size过小，很容易就会被BN层影响结果
'''
xmodel.eval()

# xmodel为载入模型
'''
将pt模型文件转为onnx模型文件--第一个参数为模型文件，第二个参数为输入模型的数据的shape
（输入模型数据的shape要在此设置，channel,width,height与模型训练时的输入数据相当,由于一般只输入一张图像测试,batch_size设置为1即可）
'''
torch.onnx.export(xmodel, x, OUT_ONNX, export_params=True, training=False, input_names=input_names,
                  output_names=out_names)
print('please run: python3 -m onnxsim test.onnx  test_sim.onnx\n')  # 对转化后的模型，运行这个命令，简化符号
print('convert done!\n')
